from typing import Dict

import numpy as np

from config.config import get_my_global_config


class FeatureBuilder(object):
    def __init__(self):
        # 缓存
        self.current_step = None
        self.traffic_light = None
        self.side_walks = None
        self.cross_walks = None
        self.all_traffic_data = None
        self.all_map_data = None
        self.env = None
        self.expert_path = None
        self._set_config()

    def _set_config(self):
        # 归一化配置
        self.max_agents = get_my_global_config()['feature']['max_agents']
        self.max_lights = get_my_global_config()['feature']['max_lights']
        self.max_sidewalks = get_my_global_config()['feature']['max_sidewalks']
        self.num_points = get_my_global_config()['feature']['num_points']
        self.max_crosswalks = get_my_global_config()['feature']['max_crosswalks']
        self.max_lanes = get_my_global_config()['feature']['max_lanes']

        self.ego_radius = get_my_global_config()['feature']['ego_radius']
        self.history_step = get_my_global_config()['base']['history_step']
        self.future_step = get_my_global_config()['base']['future_step']

        vehicle_cfg = get_my_global_config()['env']['base_config']['vehicle_config']

        self.lidar_dim = vehicle_cfg['lidar']['num_lasers']
        self.side_detector_dim = vehicle_cfg['side_detector']['num_lasers']
        self.lane_line_detector_dim = vehicle_cfg['lane_line_detector']['num_lasers']
        self.ego_dim = 6
        self.navi_dim = 22
        self.other_v_dim = vehicle_cfg['lidar']['num_others'] * 4

    def set_env(self, env):
        self.env = env

    def clear_feature(self):
        self.all_traffic_data = None
        self.all_map_data = None
        self.cross_walks = None
        self.expert_path = None
        self.side_walks = None
        self.traffic_light = None

    def _get_obs_feature(self, obs):
        """
            Ego states: [
                [Side_detector cloud_points] dim=120

                Difference of heading between ego vehicle and current lane,
                Current speed,
                Current steering,
                Throttle/brake of last frame,
                Steering of last frame,
                Yaw Rate,

                 [lane_line_detector cloud points] dim=120
                ], dim >= 9

            Navi info: [
                        Projection of distance between ego vehicle and checkpoint on ego vehicle's heading direction,
                        Projection of distance between ego vehicle and checkpoint on ego vehicle's side direction,
                        Radius of the lane whose end node is the checkpoint (0 if this lane is straight),
                        Clockwise (1) or anticlockwise (0) (0 if lane is straight),
                        Angle of the lane (0 if lane is straight)
                       ] * 2, dim = 22

            Since agent observes current lane info and next lane info, and two checkpoints exist, the dimension of navi info
            is 22.

            State observation + Navi info + 4 * closest vehicle info + Lidar points ,
            Definition of State Observation and Navi information can be found in **class StateObservation**
            Other vehicles' info: [
                                  Projection of distance between ego and another vehicle on ego vehicle's heading direction,
                                  Projection of distance between ego and another vehicle on ego vehicle's side direction,
                                  Projection of speed between ego and another vehicle on ego vehicle's heading direction,
                                  Projection of speed between ego and another vehicle on ego vehicle's side direction,
                                  ] * 4, dim = 40

            Lidar points: 240 lidar points surrounding vehicle, starting from the vehicle head in clockwise direction

            :param vehicle: BaseVehicle
            :return: observation in 120 + 6 + 120 + 22 + 40 + 240 dim
            """
        # 使用一个数组存储每个维度的偏移量，避免重复计算
        dim_offsets = [
            self.side_detector_dim,
            self.ego_dim,
            self.lane_line_detector_dim,
            self.navi_dim,
            self.other_v_dim,
            self.lidar_dim
        ]

        # 从obs中提取各个state
        sider_state = obs[:dim_offsets[0]]
        ego_state = obs[dim_offsets[0]:dim_offsets[0] + dim_offsets[1]]
        lane_state = obs[dim_offsets[0] + dim_offsets[1]:dim_offsets[0] + dim_offsets[1] + dim_offsets[2]]
        navi_state = obs[
                     dim_offsets[0] + dim_offsets[1] + dim_offsets[2]:dim_offsets[0] + dim_offsets[1] + dim_offsets[2] +
                                                                      dim_offsets[3]]
        other_v_state = obs[
                        dim_offsets[0] + dim_offsets[1] + dim_offsets[2] + dim_offsets[3]:dim_offsets[0] + dim_offsets[
                            1] + dim_offsets[2] + dim_offsets[3] + dim_offsets[4]]
        lidar_state = obs[:dim_offsets[5]]

        obs_feats = {
            "sider_state": sider_state,
            "ego_state": ego_state,
            "lane_state": lane_state,
            "navi_state": navi_state,
            "other_v_state": other_v_state,
            "lidar_state": lidar_state,
        }
        return obs_feats

    def _get_ego_features(self) -> Dict:
        agent = self.env.agent

        """获取自车特征"""
        return {
            "position": agent.position,
            "velocity": agent.velocity,
            "heading": agent.heading,
        }

    def _get_expert_future(self):
        tm = self.env.engine.traffic_manager
        ego_id = tm.sdc_scenario_id  # SDC 在 traffic_data 中的 key
        raw = tm.current_traffic_data.get(ego_id)

        if raw is None:
            # 万一没有这条 track，返回空
            return {
                "position": np.zeros((0, 2), dtype=np.float32),
                "heading": np.zeros((0,), dtype=np.float32),
                "velocity": np.zeros((0, 2), dtype=np.float32),
                "valid": np.zeros((0,), dtype=np.float32),
            }

        state = raw["state"]


        # 获取position、velocity和valid
        position = np.array(state["position"][:, :2], dtype=np.float32)  # 只取前两个维度
        velocity = np.array(state["velocity"], dtype=np.float32)
        valid = np.array(state["valid"], dtype=np.float32)

        # 计算当前步加未来步的结束位置
        end_idx = self.current_step + self.future_step +1

        # 如果越界，则使用最后一个状态填充
        if end_idx > len(position):
            # 计算缺失的步数
            missing_steps = end_idx - len(position)
            # 使用最后一个状态填充缺失的部分
            position = np.pad(position, ((0, missing_steps), (0, 0)), mode='edge')
            velocity = np.pad(velocity, ((0, missing_steps), (0, 0)), mode='edge')
            valid = np.pad(valid, (0, missing_steps), mode='constant', constant_values=0)

        # 获取从current_step到current_step + future_step的部分数据
        expert_path = {
            "position": position[self.current_step+1:end_idx, :],
            "velocity": velocity[self.current_step+1:end_idx, :],
            "valid": valid[self.current_step+1:end_idx],
        }

        return expert_path

    def get_nums_step(self):
        tm = self.env.engine.traffic_manager
        ego_id = tm.sdc_scenario_id  # SDC 在 traffic_data 中的 key
        raw = tm.current_traffic_data.get(ego_id)
        state = raw["state"]

        return len(state["position"])


    def get_feature(self, obs):
        self.current_step = self.env.episode_step
        data = {'obs_feats': self._get_obs_feature(obs), "ego": self._get_ego_features(),
                'expert_path': self._get_expert_future()}

        return data
